from multiprocessing import Queue
from loguru import logger
import gzip
import websockets
from typing import Any, Callable, Dict, Optional
from ..models.protobuf.douyin import *
from datetime import datetime
from ...utils.liveutli import extract_filename


class DouyinMessageHandler:
    def __init__(self) -> None:
        """
        初始化消息处理器
        """
        self.consecutive_trigger_count: Dict[str, int] = {}
        self.last_triggered_url: Optional[str] = None

    async def handle_message(self,
                           message: bytes,
                           url: str,
                           ws: websockets.WebSocketClientProtocol,
                           shared_queue: Any,
                           shared_list: list,
                           on_close_callback: Callable) -> None:
        """
        处理接收到的WebSocket消息
        
        Args:
            message: 接收到的二进制消息
            url: 直播间URL
            ws: WebSocket连接实例
            shared_queue: 共享队列
            shared_list: 共享列表
            on_close_callback: 连接关闭回调函数
        """
        
        package = PushFrame().parse(message)
        response = Response().parse(gzip.decompress(package.payload))

        if response.need_ack:
            await self._send_ack(ws, package, response)

        for msg in response.messages_list:
            method = msg.method
            if 'WebcastControlMessage' == method:
                if await self._handle_control_message(msg, url, ws, shared_list, on_close_callback):
                    return

            try:
                await self._process_message(msg, method, url, shared_queue)
            except Exception as e:
                logger.error(f"处理消息错误: {url} - {str(e)}")

    async def _send_ack(self,
                       ws: websockets.WebSocketClientProtocol,
                       package: PushFrame,
                       response: Response) -> None:
        """
        发送确认消息
        """
        try:
            ack = PushFrame(
                log_id=package.log_id,
                payload_type='ack',
                payload=response.internal_ext.encode('utf-8')
            ).SerializeToString()
            await ws.send(ack)
        except Exception as e:
            # logger.error(f"发送ACK失败: {str(e)}")
            pass

    async def _handle_control_message(self,
                                    msg: Message,
                                    url: str,
                                    ws: websockets.WebSocketClientProtocol,
                                    shared_list: list,
                                    on_close_callback: Callable) -> bool:
        """
        处理控制消息
        """
        message = ControlMessage().parse(msg.payload)
        if message.status == 3:
            logger.info(f"{url} 直播已结束")
            await on_close_callback(ws, url, shared_list, "直播结束")
            return True
        return False

    async def _process_message(self,
                             msg: Message,
                             method: str,
                             url: str,
                             shared_queue: Queue) -> None:
        """
        处理业务消息
        """
        message_parsers = {
            'WebcastChatMessage': ChatMessage().parse,
            'WebcastGiftMessage': GiftMessage().parse,
            'WebcastLikeMessage': LikeMessage().parse,
            'WebcastMemberMessage': MemberMessage().parse,
        }

        if method not in message_parsers:
            return

        message = message_parsers[method](msg.payload)
        await self._process_verified_message(message, url, shared_queue)

    async def _process_verified_message(self,
                                      message: Any,
                                      url: str,
                                      shared_queue: Queue) -> None:
        """
        处理已验证的消息
        """
        user_data = {
            "t_user": message.user.display_id,
            "t_name": message.user.nick_name,
            "t_home": message.user.sec_uid,
            "t_level": await self._get_user_level(message),
            "t_imgage": extract_filename(message.user.avatar_thumb.url_list_list[0]),
            "url": url,
            "day": datetime.now().strftime("%Y/%m/%d")
        }
        shared_queue.put(user_data)
        # logger.info(f"接收到消息: {shared_queue.qsize()}")

    async def _get_user_level(self, message: Any) -> int:
        """
        获取用户等级
        """
        try:
            if "荣誉等级" in message.user.badge_image_list[0].content.alternative_text:
                return message.user.badge_image_list[0].content.level
        except Exception:
            pass
        return 0 